import os
import sys
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics
import numpy as np

np.random.seed(777)
tf.random.set_seed(777)

N_EPOCHS = 1000  # 为了快速演示，这个值设置的偏小，正式代码要适当调大一些
ALPHA = 0.01
N_RNN_HIDDEN = 10

# 2.	使用keras完成字符串预测（每题7分，共42分）
# ①	数据处理
# 1)	创建字符串' if you want you like'（7分）
sentence = ' if you want you like'

# 2)	使用' if you want you lik'预测'if you want you like'（7分）
x1_str = sentence[:-1]
y1_str = sentence[1:]
N_STEPS = len(x1_str)

# 3)	对数据进行合理预处理（7分）
dict = set(list(sentence))
LEN_DICT = len(dict)
idx2char = list(dict)
char2idx = {ch: i for i, ch in enumerate(idx2char)}
x1_idx = [char2idx[ch] for ch in x1_str]
y1_idx = [char2idx[ch] for ch in y1_str]
x1_idx = np.array(x1_idx)
y1_idx = np.array(y1_idx)
x1 = np.eye(LEN_DICT)[x1_idx]
y1 = np.eye(LEN_DICT)[y1_idx]
x = tf.expand_dims(x1, axis=0)
y = tf.expand_dims(y1, axis=0)

# ②	模型创建
# 1)	模型使用两层长短期记忆（7分）
model = keras.Sequential([
    layers.LSTM(N_RNN_HIDDEN, return_sequences=True, input_shape=(N_STEPS, LEN_DICT)),
    layers.LSTM(N_RNN_HIDDEN, return_sequences=True),
    layers.TimeDistributed(layers.Dense(1)),
])
model.summary()

# 2)	进行训练，合理选择优化器，损失函数和循环次数（7分）
model.compile(
    optimizer=optimizers.Adam(learning_rate=ALPHA),
    loss=losses.categorical_crossentropy,
    metrics=[metrics.categorical_accuracy]
)
model.fit(x, y, batch_size=1, epochs=N_EPOCHS)

# ③	模型预测
# 1)	打印预测值对应的字符串（7分）
print('打印预测值对应的字符串')
h = model.predict(x)
h_idx = np.argmax(h, axis=2)
h_str = [''.join([idx2char[i] for i in row]) for row in h_idx]
print(h_str)
